import React, { useReducer, useRef } from 'react';
import { useXAgent, Sender } from '@ant-design/x';
import "./index.css"
// 消息接口定义
export interface Message {
  role: string;
  content: string;
}

export interface ResponseObject {
  model: string;
  created_at: string;
  message: Message;
  done_reason: string;
  done: boolean;
  total_duration: number;
  load_duration: number;
  prompt_eval_count: number;
  prompt_eval_duration: number;
  eval_count: number;
  eval_duration: number;
}

// 状态管理初始值
const initialState = {
  conversationHistory: [] as Message[],
  currentAssistantMessage: '',
  loading: false,
  error: '',
};

// 状态管理 reducer
function reducer(state: typeof initialState, action: any) {
  switch (action.type) {
    case 'ADD_USER_MESSAGE':
      return { ...state, conversationHistory: [...state.conversationHistory, action.payload] };
    case 'UPDATE_ASSISTANT_MESSAGE':
      return { ...state, currentAssistantMessage: action.payload };
    case 'ADD_ASSISTANT_MESSAGE':
      return {
        ...state,
        conversationHistory: [
          ...state.conversationHistory,
          { content: state.currentAssistantMessage, role: 'assistant' },
        ],
        currentAssistantMessage: '',
      };
    case 'SET_LOADING':
      return { ...state, loading: action.payload };
    case 'SET_ERROR':
      return { ...state, error: action.payload };
    default:
      return state;
  }
}

const App = () => {
  const [agent] = useXAgent({
    baseURL: 'http://localhost:11434/api/chat',
    model: 'llama3.1',
  });

  const [state, dispatch] = useReducer(reducer, initialState);
  const currentMessageRef = useRef(''); // 缓存助手消息，减少频繁更新

  const chatRequest = (text: string) => {
    dispatch({ type: 'SET_ERROR', payload: '' });
    dispatch({ type: 'SET_LOADING', payload: true });
    dispatch({ type: 'ADD_USER_MESSAGE', payload: { content: text, role: 'user' } });

    const requestInfo = {
      messages: [
        {
          content: text,
          role: 'user',
        },
      ],
      headers: { 'Content-Type': 'application/json' },
      stream: true,
    };

    const formattedMessages = requestInfo.messages.map((message) => JSON.stringify(message));

    agent.request(formattedMessages, {
      onError: (error) => {
        console.error('Error occurred:', error);
        dispatch({ type: 'SET_ERROR', payload: 'Something went wrong. Please try again.' });
        dispatch({ type: 'SET_LOADING', payload: false });
      },
      onUpdate: (message: string) => {
        try {
          const update: ResponseObject = JSON.parse(message);
          if (update.message && update.message.content) {
            currentMessageRef.current += update.message.content;
            dispatch({ type: 'UPDATE_ASSISTANT_MESSAGE', payload: currentMessageRef.current });
          }
        } catch (error) {
          console.error('Failed to parse message:', error);
        }
      },
      onSuccess: () => {
        dispatch({ type: 'ADD_ASSISTANT_MESSAGE' });
        currentMessageRef.current = '';
        dispatch({ type: 'SET_LOADING', payload: false });
      },
    });
  };

  return (
    <div className="chat-container">
      {/* 错误消息显示 */}
      {state.error && <div className="error-message">{state.error}</div>}

      {/* 对话历史显示 */}
      <ul className="chat-history">
        {state.conversationHistory.map((message, index) => (
          <li key={index} className={message.role}>
            <strong>{message.role === 'user' ? 'User:' : 'Assistant:'}</strong> {message.content}
          </li>
        ))}
      </ul>

      {/* 加载指示器 */}
      {state.loading && <div className="loading-spinner">Loading...</div>}

      {/* 输入框和发送按钮 */}
      <Sender onSubmit={chatRequest} />
    </div>
  );
};

export default App;
